#include <fftw3.h>
#include <sys/time.h>

#include "cumacros.h"

#include "auxfns.h"
#include "fftfns.h"

int sampling_pattern (int*& nzmask, cplxf *data, int *dsize)
{
	int ds0 = dsize[0], ds1 = dsize[1], ds2 = dsize[2];

	nzmask = new int [ds1*ds2];

	int nnz = 0;

	int i0_0 = ds0/2-10, i0_n = ds0/2+10;
	if (i0_0 < 0) i0_0 = 0;
	if (i0_n > ds0) i0_n = ds0;


	for (int i2 = 0; i2 < ds2; ++i2)
		for (int i1 = 0; i1 < ds1; ++i1)
		{
			int& nz = nzmask[i1 + ds1*i2];

			nz = 0;

			for (int i0 = i0_0; i0 < i0_n; ++i0)
				nz |= (abs(data[i0 + ds0*(i1 + ds1*i2)]) != 0);

			nnz += nz;
		}

	return nnz;
}

void solver_mask (float *&mask, cplxf *X, int xs0, int m, int n) 
{
	int dsize[4] = {xs0, m, n};
	int *nzmask; sampling_pattern (nzmask, X, dsize);

	int (*jr)[2] = new int [m][2], (*ic)[2] = new int [n][2];

	mask = new float [m*n];

	for (int i = 0; i < m; ++i)
	{
		jr[i][0] = n+1; jr[i][1] = -1;

		for (int j = 0; j < n; ++j)
			if (nzmask[i + m*j]){
				jr[i][0] = j;
				break;
			}

		for (int j = n-1; j >= 0; --j)
			if (nzmask[i + m*j]){
				jr[i][1] = j;
				break;
			}
	}

	for (int j = 0; j < n; ++j)
	{
		ic[j][0] = m+1; ic[j][1] = -1;

		for (int i = 0; i < m; ++i)
			if (nzmask[i + m*j]){
				ic[j][0] = i;
				break;
			}

		for (int i = m-1; i >= 0; --i)
			if (nzmask[i + m*j]){
				ic[j][1] = i;
				break;
			} 
	}

	for (int j = 0; j < n; ++j)
		for (int i = 0; i < m; ++i)
		{
			if (nzmask [i + m*j] or
			    not ((i >= ic[j][0] and i <= ic[j][1]) or
			         (j >= jr[i][0] and j <= jr[i][1])))
				mask[i + m*j] = 1.0f;
			else
				mask[i + m*j] = 0.0f;
		}

	delete [] jr; delete[] ic;
	delete [] nzmask;
}


void calib_size (int& si, int& sj, int *mask, int M, int N)
{
	int fi = 0, fj = 0;

	si=2; sj=2;

	while (!fi || !fj)
	{
		if (!fi)
		{
			int i0, j0;
			pad_indices (i0, j0, M, N, si+1, sj);

			for (int j = j0; j < j0+sj; j++)
				for  (int i = i0; i < i0+si+1; i++)
					if (mask[i + M*j] == 0){
						fi = 1;
						goto fi_out;
					}

			fi_out: if (!fi) si++;
		}

		if (!fj)
		{
			int i0, j0;
			pad_indices (i0, j0, M, N, si, sj+1);

			for (int j = j0; j < j0+sj+1; j++)
				for (int i = i0; i < i0+si; i++)
					if (mask[i + M*j] == 0){
						fj = 1;
						goto fj_out;
					}

			fj_out: if (!fj) sj++;
		}
	}
}

void
ifft_readout (cplxf *data, int *dsize)
{
	int ds0 = dsize[0], ds1 = dsize[1]*dsize[2]*dsize[3]; 

	for (int i1 = 0; i1 < ds1; ++i1)
		ifftshift (&data[i1*ds0], ds0);

	fftwf_iodim xform[1] = {{ds0, 1, 1}};
	fftwf_iodim vectr[1] = {{ds1, ds0, ds0}};

	fftwf_plan p = fftwf_plan_guru_dft (1, xform, 1, vectr,
	                                    (fftwf_complex*) data,
	                                    (fftwf_complex*) data,
	                                    FFTW_BACKWARD, FFTW_ESTIMATE);
	fftwf_execute (p);
	fftwf_destroy_plan (p);

	for (int i1 = 0; i1 < ds1; ++i1)
		fftshift (&data[i1*ds0], ds0);
}

void
fft_readout (cplxf *data, int *dsize)
{
	int ds0 = dsize[0], ds1 = dsize[1]*dsize[2]*dsize[3]; 

	for (int i1 = 0; i1 < ds1; ++i1)
		ifftshift (&data[i1*ds0], ds0);

	fftwf_iodim xform[1] = {{ds0, 1, 1}};
	fftwf_iodim vectr[1] = {{ds1, ds0, ds0}};

	fftwf_plan p = fftwf_plan_guru_dft (1, xform, 1, vectr,
	                                    (fftwf_complex*) data,
	                                    (fftwf_complex*) data,
	                                    FFTW_FORWARD, FFTW_ESTIMATE);
	fftwf_execute (p);
	fftwf_destroy_plan (p);

	for (int i1 = 0; i1 < ds1; ++i1)
		fftshift (&data[i1*ds0], ds0);
}

void
ifft_kern_readout (cplxf *&k2d, cplxf *k3d, int *ksize, int *dsize)
{
	int ds0 = dsize[0], nch = dsize[3];
	int ks0 = ksize[0], ks1 = ksize[1], ks2 = ksize[2];

	k2d = new cplxf[ds0*ks1*ks2*nch*nch];
	for (int i = 0; i < ds0*ks1*ks2*nch*nch; ++i) k2d[i] = 0.0f;

	int i0_0 = pad_index (ds0, ks0);

	for (int i1 = 0; i1 < ks1*ks2*nch*nch; ++i1){
		for (int i0 = 0; i0 < ks0; ++i0)
			k2d [i0_0+i0 + ds0*i1] = k3d [ks0-1-i0 + ks0*i1];

		ifftshift (&k2d[ds0*i1], ds0);
	}

	fftwf_iodim xform[1] = {{ds0, 1, 1}};
	fftwf_iodim hmany[1] = {{ks1*ks2*nch*nch, ds0, ds0}};

	fftwf_plan p = fftwf_plan_guru_dft (1, xform, 1, hmany,
	                                    (fftwf_complex*) k2d,
	                                    (fftwf_complex*) k2d,
	                                    FFTW_BACKWARD, FFTW_ESTIMATE);
	fftwf_execute (p);
	fftwf_destroy_plan (p);

	for (int i1 = 0; i1 < ks1*ks2*nch*nch; ++i1){
		fftshift (&k2d[ds0*i1], ds0);
	}
}

void
crop_calib_data (cplxf *&calib, int *csize, cplxf *data, int *dsize, int nr)
{
	int ds0 = dsize[0], ds1 = dsize[1], ds2 = dsize[2], nch = dsize[3]; 
	int &cs0 = csize[0], &cs1 = csize[1], &cs2 = csize[2], &cs3 = csize[3]; 

	double echo_er = -1.0;
	int echo = -1;

	for (int i0 = 0; i0 < ds0; ++i0)
	{
		double er = 0;

		for (int i1 = 0; i1 < ds1*ds2*nch; ++i1)
			er += abs2(data[i0 + ds0*i1]);

		if (er > echo_er){
			echo_er = er; echo = i0;
		}
	}

	int *nzmask; sampling_pattern (nzmask, data, dsize); 

	calib_size (cs1, cs2, nzmask, ds1, ds2);


	cs0 = nr; cs3 = nch;

	calib = new cplxf [cs0*cs1*cs2*cs3];

	int i1_0, i2_0; pad_indices (i1_0, i2_0, ds1, ds2, cs1, cs2);
	int i0_0 = echo-(cs0/2);

	for (int ch = 0; ch < nch; ++ch)
		for (int i2 = 0; i2 < cs2; ++i2)
			for (int i1 = 0; i1 < cs1; ++i1)
				for (int i0 = 0; i0 < cs0; ++i0)
					calib [i0 + cs0*(i1 + cs1*(i2 + cs2*ch))]
					  = data [i0_0+i0 + ds0*(i1_0+i1 + ds1*(i2_0+i2 + ds2*ch))];

	delete [] nzmask;
}

/*
int detect_gpus (void)
{
	int ngpus;

	cuda(GetDeviceCount (&ngpus));

	if (ngpus == 1)
	{
		cudaDeviceProp prop;
		cuda(GetDeviceProperties(&prop, 0));

		if (prop.major == 9999 and prop.minor == 9999) ngpus = 0;
	}

	return ngpus;
}
*/

/*double timestamp()
{
	struct timeval tv;
	gettimeofday (&tv, 0);
	return tv.tv_sec + 1e-6*tv.tv_usec;
}*/
